import numpy as np
import os
from structAIDetect.simdatagen.getfiles import get_files_by_extension
import matplotlib.pyplot as plt

THICKNESS = 0.5


def data_norm(datapath, nodeid_absmax, corr_absmax, shape_abs_max, freqs_abs_max, savepath):
    modaldata_dir = os.path.join(datapath)
    filesnames = get_files_by_extension(modaldata_dir, '.npz')
    for filename in filesnames:
        data = np.load(os.path.join(datapath, filename))
        cellinfo = data['cellinfo']
        cellstate = data['cellstate']
        freqs = data['freqs']
        cellinfo[:,:,0] = cellinfo[:,:,0] / nodeid_absmax
        cellinfo[:,:,1:4] = cellinfo[:,:,1:4] / corr_absmax
        for shape in range(5):
            cellinfo[:,:,4+shape*3:7+shape*3] = cellinfo[:,:,4+shape*3:7+shape*3] / shape_abs_max[shape]
        freqs = freqs / freqs_abs_max
        np.savez(os.path.join(savepath, filename), cellinfo=cellinfo, cellstate=cellstate, freqs=freqs)


def node_grid_mesh(cellinfo):
    cellinfo = cellinfo.reshape(cellinfo.shape[0] * cellinfo.shape[1], -1)
    unique_value, indice = np.unique(cellinfo[:, 0], return_index=True)
    nodeinfo = cellinfo[indice]
    nodeinfo = nodeinfo[nodeinfo[:, 3] > THICKNESS / 2]
    unique_x = np.unique(nodeinfo[:, 1])
    unique_y = np.unique(nodeinfo[:, 2])
    rows = len(unique_y)
    cols = len(unique_x)
    sorted_indice = np.lexsort((nodeinfo[:, 1], nodeinfo[:, 2]))
    sorted_node = nodeinfo[sorted_indice]
    nodegrid = sorted_node.reshape(rows, cols, -1)

    return nodegrid, rows, cols


def disfieldgen(datapath, savepath):
    """
    取平板上表面的位移数据，将其转换成以为二维位移场形式的ndarray并保存，同时过滤到其中非z轴方向的模态振型。

    参数
    --------
    datapath : str
        源文件读取路径，源文件中保存了单元形式的振动模态数据，这些数据通过cellinfo_generate()函数生成。
    savepath : str
        转换为位移场后，位移场的文件保存路径。

    文件输出
    --------
    nodegrid : ndarray
        保存位移场数据，nodegrid.shape = (ygrid_num, xgrid_num, 19)
        第1维每一位对应一个y轴方向的坐标。
        第2维每一位对应一个x轴方向的坐标。
        第3位
            0位是节点编号；
            1~3位是节点坐标；
            4~6位是1阶模态振型位移；
            7~9位是2阶模态振型位移；
            ...
            16~18位是5阶模态振型位移。
    freqs : ndarray

    """

    modaldata_dir = os.path.join(datapath)
    filesnames = get_files_by_extension(modaldata_dir, '.npz')
    idx = 0
    file_indice = []
    for filename in filesnames:
        try:
            part1 = filename.split("sim")[1]
            file_indice.append(int(part1.split(".npz")[0]))
            idx += 1
        except IndexError:
            print("未找到匹配的模式")
    file_indice = np.argsort(file_indice).tolist()
    filesnames = [filesnames[i] for i in file_indice]
    for filename in filesnames[0:]:
        data = np.load(os.path.join(datapath, filename))
        cellinfo = data['cellinfo']
        cellstate = data['cellstate']
        freqs = data['freqs']
        nodegrid, rows, cols = node_grid_mesh(cellinfo)
        rows -= 1
        cols -= 1
        cellgrid = cell_grid_mesh(cellinfo, cellstate, rows, cols)
        # print(freqs[0])
        print(sum(~cellstate), freqs)
        # for mode in range(5):
        #     gridplot(nodegrid, mode)
        nodegrid, freqs = displace_pick_by_axis(nodegrid, freqs, 2, 4)
        
        np.savez(os.path.join(savepath, 'Norm_' + filename), nodegrid=nodegrid, cellgrid=cellgrid, freqs=freqs)


def cell_grid_mesh(cellinfo, cellstate, rows, cols):
    cellcoor = cellinfo[:, :, 1:3]
    cellcoor = np.average(cellcoor, axis=1)
    sorted_indice = np.lexsort((cellcoor[:, 0], cellcoor[:, 1]))
    cellstate = cellstate[sorted_indice]
    cellgrid = cellstate.reshape(rows, cols, -1)

    return cellgrid


def displace_pick_by_axis(nodegrid, freqs, axis, mode_num_reserve):
    """
    如果某一阶模态的振动数据中，axis方向的振型参与度过低，则去除这一阶模态。

    输入
    ---------
    axis : int
        0表示x轴方向，1表示y轴方向，2表示z轴方向
    """

    if mode_num_reserve == 0:
        raise ValueError('mode_num_reserve cannot be 0')
    
    if axis > 2:
        raise ValueError('axis cannot > 2')
    
    total_mode_num = (nodegrid.shape[2] - 4) // 3
    mode_popout = []
    mode_shape_popout = []
    for mode in range(total_mode_num):
        displacement_x = nodegrid[:, :, 4 + mode]
        displacement_x = displacement_x.reshape(-1, 1)
        displacement_y = nodegrid[:, :, 4 + mode + total_mode_num]
        displacement_y = displacement_y.reshape(-1, 1)
        displacement_z = nodegrid[:, :, 4 + mode + total_mode_num * 2]
        displacement_z = displacement_z.reshape(-1, 1)
        displacement_axis = nodegrid[:, :, 4 + mode + total_mode_num * axis]
        displacement_axis = displacement_axis.reshape(-1, 1)
        norm_displacement = np.linalg.norm(displacement_x, ord=2) + np.linalg.norm(displacement_y, ord=2) + np.linalg.norm(displacement_z, ord=2)
        norm_displacement_axis = np.linalg.norm(displacement_axis, ord=2)
        if norm_displacement_axis < 0.1 * norm_displacement / 3:
            mode_popout.append(mode)
            mode_shape_popout.append(4 + mode)
            mode_shape_popout.append(4 + mode + total_mode_num)
            mode_shape_popout.append(4 + mode + total_mode_num * 2)

    if not mode_popout:
        for mode in range(mode_num_reserve, total_mode_num):
            mode_popout.append(mode)
            mode_shape_popout.append(4 + mode)
            mode_shape_popout.append(4 + mode + total_mode_num)
            mode_shape_popout.append(4 + mode + total_mode_num * 2)
    elif len(mode_popout) > total_mode_num - mode_num_reserve:
        raise ValueError('Invalid mode larger than expected')

    nodegrid = np.delete(nodegrid, mode_shape_popout, axis=axis)
    freqs = np.delete(freqs, mode_popout)
    
    print(mode_popout)

    return nodegrid, freqs

def gridplot(nodegrid, mode):
    plt.figure()
    ax = plt.subplot(projection='3d')
    mode_order = mode + 1
    amp = 50 / np.max(np.abs(nodegrid[:, :, 4 + (mode_order - 1)::5]))
    for idx in range(nodegrid.shape[0]):
        curve_x = nodegrid[idx, :, :]
        ax.plot(curve_x[:, 1] + curve_x[:, 4 + (mode_order - 1)] * amp, 
                curve_x[:, 2] + curve_x[:, 4 + (mode_order - 1) + 5] * amp,
                curve_x[:, 3] + curve_x[:, 4 + (mode_order - 1) + 10] * amp)
    ax.axis('equal')
    plt.show()



